import torch
# matplotlib.pyplot 是 Python 最著名的 2D 绘图库，专为数据可视化设计。
import matplotlib.pyplot as plt

# 生成输入数据
x = torch.linspace(-5, 5, 500, requires_grad=True)

# 定义激活函数字典
activations = {
    "Sigmoid": lambda x: 1 / (1 + torch.exp(-x)),
    "Tanh": torch.tanh,
    "ReLU": torch.relu,
    "LeakyReLU (α=0.1)": lambda x: torch.where(x > 0, x, 0.1 * x)
}

# 绘制每个激活函数及其导数
for name, func in activations.items():
    # 计算函数值和导数
    y = func(x)
    y.sum().backward()  # 自动求导
    dy = x.grad.clone()
    x.grad.zero_()  # 清零梯度

    # 创建新的figure
    plt.figure(figsize=(10, 4))

    # 绘制函数曲线
    plt.subplot(1, 2, 1)
    plt.plot(x.detach().numpy(), y.detach().numpy(), 'b-')
    plt.title(f"{name}函数", fontsize=12)
    plt.xlabel("Input")
    plt.ylabel("Output")
    plt.grid(True)

    # 绘制导数曲线
    plt.subplot(1, 2, 2)
    if name in ["ReLU", "LeakyReLU (α=0.1)"]:
        plt.step(x.detach().numpy(), dy.detach().numpy(), 'r-')
    else:
        plt.plot(x.detach().numpy(), dy.detach().numpy(), 'r-')
    plt.title(f"{name}导数", fontsize=12)
    plt.xlabel("Input")
    plt.ylabel("Derivative")
    plt.grid(True)

    plt.tight_layout()
    plt.savefig(f'{name}.png', dpi=150)  # 保存为单独文件
    plt.show()